//
// Created by 86180 on 2025/1/18.
//

#ifndef AVLTREE_AVLMAP_H
#define AVLTREE_AVLMAP_H

#include <algorithm>
#include <vector>
#include <string>
#include <functional>
template <typename Key,typename Value>
struct AVLNode{
    Key key;
    Value value;
    int height;
    AVLNode* left;
    AVLNode* right;
    AVLNode(const Key& key,const Value& value):key(key),value(value),height(1),left(nullptr),right(nullptr){}
};

template <typename Key,typename Value>
int getHeight(AVLNode<Key, Value> *node) {
    if(node == nullptr)return 0;
    return node->height;
}

template <typename Key,typename Value>
int getBalance(AVLNode<Key,Value>* node) {
    if(node == nullptr)return 0;
    return getHeight(node->left) - getHeight(node->right);
}

/*
     y                             x
    / \                           / \
   x   T3            ==>          z   y
  / \                               / \
 z   T2                            T2  T3
 * */

template <typename Key,typename Value>
AVLNode<Key,Value>* rightRotate(AVLNode<Key,Value>* y){
    AVLNode<Key,Value>* x = y->left;

    y->left = x->right;
    x->right = y;

    y->height = std::max(getHeight(y->left), getHeight(y->right)) + 1;
    x->height = std::max(getHeight(x->left), getHeight(x->right)) + 1;

    return x;
}
/*
    x                             y
   / \                           / \
  T1  y           ==>            x   z
     / \                       / \
    T2  z                     T1  T2
 */
template <typename Key,typename Value>
AVLNode<Key,Value>* leftRotate(AVLNode<Key,Value>* x){
    AVLNode<Key,Value>* y = x->right;

    x->right = y->left;
    y->left = x;

    y->height = std::max(getHeight(y->left), getHeight(y->right)) + 1;
    x->height = std::max(getHeight(x->left), getHeight(x->right)) + 1;

    return y;
}

template <typename Key, typename Value>
AVLNode<Key, Value>* insertNode(AVLNode<Key, Value>* node, const Key& key, const Value& value){
    if(node == nullptr){
        return new AVLNode<Key,Value>(key,value);
    }
    if(node->key>key){
        node->left = insertNode(node->left,key,value);
    }else if(node->key<key){
        node->right = insertNode(node->right,key,value);
    }else{
        node->value = value;
        return node;
    }

    node->height = 1+std::max(getHeight(node->left), getHeight(node->right));
    //根据平衡因子判断左右旋
    int balance = getBalance(node);

    if (balance > 1 && key < node->left->key)
        return rightRotate(node);

    // 右右情况
    if (balance < -1 && key > node->right->key)
        return leftRotate(node);

    // 左右情况
    if (balance > 1 && key > node->left->key) {
        node->left = leftRotate(node->left);
        return rightRotate(node);
    }

    // 右左情况
    if (balance < -1 && key < node->right->key) {
        node->right = rightRotate(node->right);
        return leftRotate(node);
    }
    return node;
}

template <typename Key, typename Value>
Value*  searchNode(AVLNode<Key,Value>* node ,const Key& key){
    if (node == nullptr)
        return nullptr;

    if (key == node->key)
        return &(node->value);
    else if (key < node->key)
        return searchNode(node->left, key);
    else
        return searchNode(node->right, key);
}

template <typename Key, typename Value> //中序遍历的后一个节点
AVLNode<Key, Value>*  getMinimum(AVLNode<Key,Value>* node){
    auto* current = node;
    while(current->left != nullptr){
        current = current->left;
    }
    return current;
}

template <typename Key, typename Value>
AVLNode<Key, Value>*  deleteNode(AVLNode<Key, Value>* root, const Key& key){
    if(root == nullptr)return nullptr;
    if(root->key<key){
        root->right = deleteNode(root->right,key);
    }else if(root->key>key){
        root->left = deleteNode(root->left,key);
    }else{
        //status : 删除节点是叶子节点、有一个子节点或有两个子节点
        if((root->left == nullptr) || (root->right == nullptr)){
            auto* temp = root->right== nullptr ? root->left :root->right;

            if (temp == nullptr) {
                temp = root;
                root = nullptr;
            }
            else // 一个子节点
                *root = *temp; // 复制内容

            delete temp;
        }else{
            auto* successor = getMinimum(root->right);
            root->key = successor->key;
            root->value = successor->value;
            root->right = deleteNode(root->right,successor->key);
        }
    }
    if(root == nullptr){
        return root;
    }
    root->height = 1 + std::max(getHeight(root->left), getHeight(root->right));
    //根据平衡因子判断左右旋
    int balance = getBalance(root);
    // 左左情况
    if (balance > 1 && getBalance(root->left) >= 0)
        return rightRotate(root);

    // 左右情况
    if (balance > 1 && getBalance(root->left) < 0) {
        root->left = leftRotate(root->left);
        return rightRotate(root);
    }

    // 右右情况
    if (balance < -1 && getBalance(root->right) <= 0)
        return leftRotate(root);

    // 右左情况
    if (balance < -1 && getBalance(root->right) > 0) {
        root->right = rightRotate(root->right);
        return leftRotate(root);
    }
    return root;
}



template <typename Key,typename Value>
class AVLMap {
public:
    AVLMap():root(nullptr){}

    void put(const Key& key,const Value& value){
        root = insertNode(root,key,value);
    }

    Value* get(const Key& key){
        return searchNode(root,key);
    }

    void remove(const Key& key){
        root = deleteNode(root, key);
    }

    std::vector<std::pair<Key,Value>> inorder(){
        std::vector<std::pair<Key,Value>> res;
        inorderHelper(root,res);
        return res;
    }

    ~AVLMap(){
        std::function<void(AVLNode<Key, Value>*)> destroy = [&](AVLNode<Key, Value>* node) {
            if (node) {
                destroy(node->left);
                destroy(node->right);
                delete node;
            }
        };
        destroy(root);
        //deleteroot(root);
    }

private:
    AVLNode<Key, Value>* root;
    void inorderHelper(AVLNode<Key,Value>* node,std::vector<std::pair<Key,Value>>& res) const {
        if(node != nullptr){
            inorderHelper(node->left,res);
            res.emplace_back(node->key,node->value);
            inorderHelper(node->right,res);
        }
    }
    void deleteroot(AVLNode<Key,Value>* root){
        if(root == nullptr) return;
        if(root->left != nullptr){
            deleteroot(root->left);
        }
        if(root->right != nullptr){
            deleteroot(root->right);
        }
        delete root;
    }
};


#endif //AVLTREE_AVLMAP_H
